"""
@Time: 2021/1/25 下午 4:39
@Author: jinzhuan
@File: ner_cn_toolkit.py
@Desc: 
"""
import torch
from cognlp import *
torch.cuda.set_device(4)


ner_toolkit = NerCNToolkit(
    model_path='../../../cognlp/data/ner/MSRA/toolkit/model.pt',
    vocabulary_path='../../../cognlp/data/ner/MSRA/toolkit/vocabulary.txt',
    device=torch.device("cuda"),
    device_ids=[4],
    max_seq_length=256)
ner_result = ner_toolkit.run('作者钱钟书是自动化所的老师')
print(ner_result)
